#coding=utf-8

'''
Ising模型哈密顿量：
# \hat{H} = - J \sum_{i,j相邻} s_i s_j - \mu B \sum_i s_i,
体系处于正则系综，温度为 T
'''

import numpy as np
import matplotlib.pyplot as plt

# 二维格点：L*L， 第 i 个格点的 4 个邻居的自旋总和
def Sneighbours( s, L, i):
    s_up = s[i - L] if i >= L else s[i + L*L - L]  # 上邻居
    s_down = s[i + L] if i < L*L-L else s[i - (L*L-L)]  # 下邻居
    s_left = s[i - 1] if i % L != 0 else s[i + L -1]   # 左邻居
    s_right = s[i + 1] if (i+1) % L != 0 else s[i - (L-1)]   # 右邻居
    #print( s_up, s_down, s_left, s_right )
    return s_up + s_down + s_left + s_right

'''
# test Sneighbours 
for i in range(L):
    print( [s[i*L+j] for j in range(L)] )
sum=0
for i in range(L*L):
    print("Sneighbours = ", Sneighbours(s, L, i))
    sum += Sneighbours(s, L, i)
print("sum=", sum/4, np.sum(s))
'''
# Ising模型哈密顿量：
# \hat{H} = - J \sum_{i,j相邻} s_i s_j -  B \sum_i s_i,
# Delta E = 2 J * s_i * S_{邻居} + 2 B s_i
def DeltaE( s, L, i, J, B ):
    return 2 * J * s[i] * Sneighbours(s, L, i) + 2 * B * s[i]

def Energy( s, L, J, B ):
    E = 0
    for i in range(L*L):
        E += - J * s[i] * Sneighbours(s, L, i) - B * s[i]
    return E/2

# if Ef < Ei, A = 1
# else if Ef >= Ei, A = exp{ -(Ef - Ei)/kbT }
def A( s, L, i, J, B, kbT ):
    x = DeltaE(s, L, i, J, B)
    #print("Delta E = ", x)
    if x <= 0:
        return 1
    else:
        #print( np.e ** ( -x/kbT ) )
        return np.e ** ( -x/kbT )

# flip s[i] if accepted
def tryflip( s, L, i, J, B, kbT ):
    a = A( s, L, i, J, B, kbT )
    if np.random.random() < a :
        s[i] *= -1
        #print("accepted")

# total magnetic moment
def Ising( nMarkov, nstep, L, J, B, kT):
    M = []
    E = []
    for i in range(nMarkov):
        s = [2 * np.random.randint(0, 2) - 1 for i in range(L * L)]
        for j in range(nstep):
            #print("s=",s)
            for k in range(L * L):
                tryflip(s, L, k, J, B, kT )
            #print("s=",s)
            #exit(1)
        #print("M=", sum(s))
        M.append( abs(np.average(s)) )
        E.append( Energy(s, L, J, B)/(L*L) )
        #print(i, sum_M)
    return E, M

nMarkov = 1000; nstep = 10000; L = 30; J = 1; kT = 5; B = 0

ave_M = []; ave_E = [];
kTdata = np.arange(0.1, 5, 0.1)
for kT in kTdata:
    print("kT = ", kT)
    x = Ising(nMarkov, nstep, L, J, B, kT)
    print("ave_M = ", np.average(x[1]) )
    ave_E.append( np.average(x[0]) )
    ave_M.append( np.average(x[1]) )
print("kTdata = ", kTdata)
print("ave_M = ", ave_M)
plt.plot(kTdata, ave_M, "ko-")
#plt.plot(kTdata, ave_E, "r*--")
plt.xlabel("kT/J"); plt.ylabel("$<|s_z|>$")
plt.savefig("../../../doc/marp-slides/Ising2d.png")
plt.show()